/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
 * Description: mc2 graph adaptor
 * Author: lilin
 * Create: 2023-11-21
 */

#include "hcom_graph_mc2.h"
#include "register/hidden_inputs_func_registry.h"
#include "register/op_tiling_info.h"
#include "../op_base/src/op_base.h"
#include "external/hcom/hcom_topo_info.h"
#include "hcom.h"
#include "hcom_op_utils.h"

namespace hccl {
static constexpr u32 NEW_TILING_VERSION = 2U;
static const std::unordered_map<std::string, std::vector<std::string>> OP_TYPE_TO_GROUP_NAMES = {
    {"MatmulAllReduce", {"group"}},
    {"MatmulAllReduceAddRmsNorm", {"group"}},
    {"InplaceMatmulAllReduceAddRmsNorm", {"group"}},
    {"MatmulReduceScatter", {"group"}},
    {"AllGatherMatmul", {"group"}},
    {"BatchMatMulReduceScatterAlltoAll", {"group_ep", "group_tp"}},
    {"AlltoAllAllGatherBatchMatMul", {"group_ep", "group_tp"}},
    {"MoeDistributeDispatch", {"group_ep", "group_tp"}},
    {"MoeDistributeCombine", {"group_ep", "group_tp"}}
};

bool HcomGetGroupsByOpDesc(const ge::OpDescPtr &opdesc, std::vector<std::string> &groups)
{
    const auto iter = OP_TYPE_TO_GROUP_NAMES.find(opdesc->GetType());
    if (iter == OP_TYPE_TO_GROUP_NAMES.end() || iter->second.empty()) {
        HCCL_ERROR("Unsupported op %s when getting groups.", opdesc->GetName().c_str());
        return false;
    }

    std::string group;
    for (const auto groupName: iter->second) {
        CHK_PRT_RET(!ge::AttrUtils::GetStr(opdesc, groupName, group),
            HCCL_ERROR("Failed to get group %s of op %s.", groupName.c_str(), opdesc->GetName().c_str()),
            false);
        HCCL_INFO("Get group %s:%s of op %s.", groupName.c_str(), group.c_str(), opdesc->GetName().c_str());
        groups.push_back(group);
    }

    return true;
}

u32 HcomGetTilingVersionByOpDesc(const ge::OpDescPtr &opdesc, std::string &tilingData)
{
    const auto tiling = opdesc->GetExtAttr<std::shared_ptr<optiling::utils::OpRunInfo>>(ge::ATTR_NAME_OP_RUN_INFO);
    if (tiling == nullptr || *tiling == nullptr) {
        HCCL_WARNING("Failed to get tiling info of op %s.", opdesc->GetName().c_str());
        return 0U;
    }

    tilingData = (*tiling)->GetAllTilingData().str();
    const size_t tilingSize = tilingData.size();
    u32 expectedSize = sizeof(u32) + sizeof(u32);
    if (tilingSize < expectedSize) {
        HCCL_WARNING("Invalid tiling size(%zu) of op %s.", tilingSize, opdesc->GetName().c_str());
        return 0U;
    }

    return *reinterpret_cast<const u32 *>(tilingData.c_str());
}

rtStream_t HcomGetStreamByOpDesc(const ge::OpDescPtr &opdesc)
{
    const auto rtList = opdesc->TryGetExtAttr<std::shared_ptr<std::vector<void *>>>("_rt_resource_list", nullptr);
    if (rtList == nullptr) {
        HCCL_ERROR("Failed to get rt list of op %s.", opdesc->GetName().c_str());
        return nullptr;
    }

    if (rtList->size() != 1) {
        HCCL_ERROR("Invalid rt list size(%zu) of op %s.", rtList->size(), opdesc->GetName().c_str());
        return nullptr;
    }

    return (*rtList)[0UL];
}

void *HcomGetContext(const rtStream_t stream, const void *tilingData, const char *groupName)
{
    std::shared_ptr<hccl::hcclComm> hcclComm;
    HcclResult ret = HcomGetCommByGroup(groupName, hcclComm);
    CHK_PRT_RET(ret != HCCL_SUCCESS || hcclComm == nullptr,
        HCCL_ERROR("[Get][CommHandle]Failed to get hcom commHandle, group[%s] errNo[0x%016llx].",
        groupName, HCCL_ERROR_CODE(ret)), nullptr);
    HcclComm com = static_cast<HcclComm>(hcclComm.get());

    void *context = nullptr;
    if (tilingData != nullptr) {
        ret = HcclAllocComResourceByTiling(com, stream, const_cast<void *>(tilingData), &context);
    } else {
        uint64_t streamMode = 0;
        CHK_PRT_RET(hrtStreamGetMode(stream, &streamMode) != HCCL_SUCCESS,
            HCCL_ERROR("Failed to get stream mode."), nullptr);
        ret = HcclCreateComResourceByComm(com, streamMode, true, &context);
    }
    CHK_PRT_RET(ret != HCCL_SUCCESS || context == nullptr,
        HCCL_ERROR("Failed to create ComResource by tiling, errNo[0x%016llx].", HCCL_ERROR_CODE(ret)),
        nullptr);

    ge::HcomTopoInfo::TopoInfo topoInfo;
    CHK_PRT_RET(!ge::HcomTopoInfo::Instance().TryGetGroupTopoInfo(groupName, topoInfo),
        HCCL_ERROR("Failed to get topo info for group %s.", groupName),
        nullptr);

    rtStream_t opstream;
    ret = hcclComm->GetAicpuOpStreamNotify(&opstream, 1, &topoInfo.notify_handle);
    CHK_PRT_RET(ret != HCCL_SUCCESS,
        HCCL_ERROR("Failed to get Aicpu Op Stream Notify, errNo[0x%016llx].", HCCL_ERROR_CODE(ret)),
        nullptr);

    HCCL_INFO("Set notify %p to group %s.", topoInfo.notify_handle, groupName);
    ge::HcomTopoInfo::Instance().SetGroupTopoInfo(groupName, topoInfo);

    HCCL_INFO("Create context(%s) for group %s successfully, address %p.", (tilingData != nullptr ? "v2" : "v1"),
        groupName, context);
    return context;
}

REG_HIDDEN_INPUTS_FUNC(ge::HiddenInputsType::HCOM, HcomCreateComResource);
ge::graphStatus HcomCreateComResource(const ge::OpDescPtr &opdesc, std::vector<void *> &contexts)
{
    HCCL_INFO("[HcomCreateComResource]Hcomm create com resource of op %s.", opdesc->GetName().c_str());

    std::vector<std::string> groups{};
    if (!HcomGetGroupsByOpDesc(opdesc, groups) || groups.empty()) {
        HCCL_ERROR("Failed to get groups of op %s", opdesc->GetName().c_str());
        return ge::GRAPH_FAILED;
    }

    const rtStream_t stream = HcomGetStreamByOpDesc(opdesc);
    if (stream == nullptr) {
        HCCL_ERROR("Failed to get attached stream of op %s", opdesc->GetName().c_str());
        return ge::GRAPH_FAILED;
    }

    std::string tilingData;
    const u32 version = HcomGetTilingVersionByOpDesc(opdesc, tilingData);
    HCCL_INFO("Tiling version of op %s is %u.", opdesc->GetName().c_str(), version);

    for (const auto group: groups) {
        void *context = HcomGetContext(stream, reinterpret_cast<const void *>(tilingData.c_str()), group.c_str());
        CHK_PRT_RET(context == nullptr, HCCL_ERROR("Failed to create ComResource by tiling."),
            ge::GRAPH_FAILED);
        contexts.push_back(context);
    }

    return ge::GRAPH_SUCCESS;
}
}